# -*- coding: utf-8 -*-
"""
主要功能：
1. 检查标注文件中的box是否符合规则／检查object数目是否为0
2. 删除xml和图片文件文件名中的‘空格’
3. 检查xml文件是否有对应的图片与之匹配，若没有，则删除当前xml文件
4. 检查图片文件是否有xml标签文件与之匹配，若没有，则将当前图片移动到‘no_label_imgs’目录下

用法：
python check_xml.py --root_dir='标准VOC目录'
或者
修改源文件，指定VOC目录

2020-06-01
5. 检查xml文件中width和height字段是否为0，若为0则打开对应图片读取图片宽高，并更新到xml文件中
"""
import os
import shutil
import sys
import argparse
import cv2

try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET


def strip_and_rename(root_dir, dir):

    file_dir_path = os.path.join(root_dir, dir)
    renamed_num = 0
    for file in os.listdir(file_dir_path):
        old_img_name = os.path.join(file_dir_path, file)
        new_img_name = old_img_name
        if len(file.split(" ")) >= 2:
            stripped_file_name = ''.join(file.split())
            new_img_name = os.path.join(file_dir_path, stripped_file_name)
            # if len(img_file.split(" ")) > 1:
            #     num_imgs += 1
            #     old_img_name = os.path.join(img_dir, img_file)
            #     new_img_name = os.path.join(img_dir, img_file.replace(" ", "_"))
            #     os.rename(old_img_name, new_img_name)
            os.rename(old_img_name, new_img_name)
            renamed_num += 1
    print("renamed {} num: {}".format(dir, renamed_num))
    print("*" * 80)
    print()


def check_xml(root_dir, rename_files = False):
    xml_dir = os.path.join(root_dir, "Annotations")
    img_dir = os.path.join(root_dir, "JPEGImages")
    no_label_images_dir = os.path.join(root_dir, "no_label_imgs")
    if not os.path.exists(no_label_images_dir):
        os.makedirs(no_label_images_dir)
    xml_num = len(os.listdir(xml_dir))
    img_exts = [".jpg", ".jpeg", ".JPG", ".JPEG", ".PNG", ".png"]

    strip_and_rename(root_dir, "JPEGImages")
    strip_and_rename(root_dir, "Annotations")

    xml_removed = 0
    img_removed = 0

    for img_file in os.listdir(img_dir):
        old_img_name = os.path.join(img_dir, img_file)
        base_name, _ = os.path.splitext(old_img_name)
        xml_file_path = os.path.join(img_dir, base_name+".xml").replace("JPEGImages", "Annotations")

        if not os.path.exists(xml_file_path):
            shutil.move(old_img_name, os.path.join(no_label_images_dir, img_file))
            # os.remove(old_img_name)
            img_removed += 1
    print("total removed image: {}".format(img_removed))

    # remove 'space' in file name and remove file if no image file exists
    for xml_file in os.listdir(xml_dir):

        # remove 'space' in xml file in field 'filename'
        xml_path = os.path.join(xml_dir, xml_file)
        tree = ET.parse(xml_path)
        root = tree.getroot()

        full_img_basename, _ = os.path.splitext(xml_path)
        full_img_basename = full_img_basename.replace("Annotations", "JPEGImages")

        found_image = False
        for ext in img_exts:
            full_img_path = full_img_basename + ext
            if os.path.exists(full_img_path):
                found_image = True
                root.find('filename').text = full_img_path.split("/")[-1]
                tree.write(os.path.join(xml_dir, xml_file))
                break
            else:
                continue
        if not found_image:
            os.remove(xml_path)
            xml_removed += 1
            print("removed file {}, total removed XML: {}".format(xml_path, xml_removed))

    num = 0
    modified_num = 0
    for file in os.listdir(xml_dir):
        xml_file_full_path = os.path.join(xml_dir, file)
        tree = ET.parse(xml_file_full_path)
        root = tree.getroot()
        object_num = 0

        chg = 0
        img_shape = root.find('size')
        img_w = int(img_shape.find('width').text)
        img_h = int(img_shape.find('height').text)

        if img_w == 0 or img_h == 0:
            img_file_name = root.find('filename').text
            # xml_base_name = xml_file_full_path.split(".xml")[0].split("/")[-1]
            img_full_path = os.path.join(xml_dir.replace("Annotations", "JPEGImages"), img_file_name)
            # img_full_path = img_full_path.replace("Annotations", "JPEGImages")
            image = cv2.imread(img_full_path)
            image_h_from_img, image_w_from_img, _ = image.shape
            root.find('size').find('width').text = str(image_w_from_img)
            root.find('size').find('height').text = str(image_h_from_img)
            img_w = image_w_from_img
            img_h = image_h_from_img
            tree.write(xml_file_full_path)
            print('saved xml: {}, width: {}, height: {}'.format(file, img_w, img_h))

        for obj in root.iter('object'):
            object_num = object_num + 1
            xml_box = obj.find('bndbox')
            # in case that the coordinates are float type
            left = round(float(xml_box.find('xmin').text))
            top = round(float(xml_box.find('ymin').text))
            right = round(float(xml_box.find('xmax').text))
            bottom = round(float(xml_box.find('ymax').text))
            if left < 0:
                # print(xml_file_full_path.split(os.path.sep)[-1], 'xmin < 0')
                xml_box.find('xmin').text = str(1)
                chg = 1
            if left >= img_w:
                # print(xml_file_full_path.split(os.path.sep)[-1], 'xmin >= img_w')
                xml_box.find('xmin').text = str(img_w - 2)
                chg = 1
            if right < 0:
                # print(xml_file_full_path.split(os.path.sep)[-1], 'xmax < 0')
                xml_box.find('xmax').text = str(1)
                chg = 1
            if right >= img_w:
                # print(xml_file_full_path.split(os.path.sep)[-1], 'xmax >= img_w')
                xml_box.find('xmax').text = str(img_w - 2)
                chg = 1
            if top < 0:
                # print(xml_file_full_path.split(os.path.sep)[-1], 'ymin < 0')
                xml_box.find('ymin').text = str(1)
                chg = 1
            if top >= img_h:
                # print(xml_file_full_path.split(os.path.sep)[-1], 'ymin >= img_h')
                xml_box.find('ymin').text = str(img_h - 2)
                chg = 1
            if bottom < 0:
                # print(xml_file_full_path.split(os.path.sep)[-1], 'ymax < 0')
                xml_box.find('ymax').text = str(1)
                chg = 1
            if bottom >= img_h:
                # print(xml_file_full_path.split(os.path.sep)[-1], 'ymax > img_h')
                xml_box.find('ymax').text = str(img_h - 2)
                chg = 1
            if left >= right:
                print(xml_file_full_path.split(os.path.sep)[-1], 'xmin >= xmax')

            if top >= bottom:
                print(xml_file_full_path.split(os.path.sep)[-1], 'ymin >= ymax')
        if object_num == 0:
            print(xml_file_full_path.split(os.path.sep)[-1], 'has no object!')
        if chg == 1:
            tree.write(xml_file_full_path)
        modified_num += chg
        num = num + 1
        if num % 1000 == 0:
            print('completed: {}, total: {}'.format(num, xml_num))
    print('completed all xml files: {}, modified: {}'.format(num, modified_num))


if __name__ == '__main__':

    PARSER = argparse.ArgumentParser(description="arguments")
    PARSER.add_argument("--root_dir", type=str,
                        default="/zongshiban_data/2/ccj/dataset/plane_ship_tank_VOC",
                        help="VOC dataset root directory")
    ARGS = PARSER.parse_args()
    ROOT_DIR = ARGS.root_dir
    rename = False
    check_xml(ROOT_DIR, rename)
